#include <kuai/sim/md.hpp>

namespace kuai {

	void MolecularDynamicControl::start(SimulationEnvironment& env, EnergySet& eset) {
		if (env.isMainThread()) {
			rmasses.resize(eset.nAtoms);
		}
		env.wait();
		std::pair<size_t, size_t> range = env.range(eset.nAtoms);
		for (size_t i = range.first; i != range.second; ++i) {
			rmasses[i] = 1/masses[i];
		}

		env.wait();
	}

	void LeapFrogControl::act(SimulationEnvironment& env, EnergySet& eset, size_t istep) {
		RealNumber halfstep = timestep / 2;
		std::pair<size_t, size_t> range = env.range(eset.nAtoms);
		for (size_t i = range.first; i != range.second; ++i) {
			RealNumber v = halfstep * rmasses[i];
			speeds[i].x -= v * grediants[i].x;
			speeds[i].y -= v * grediants[i].y;
			speeds[i].z -= v * grediants[i].z;
			coords[i].x += timestep * speeds[i].x;
			coords[i].y += timestep * speeds[i].y;
			coords[i].z += timestep * speeds[i].z;
		}
		env.wait();
		eset.getEG(env, energies, grediants);
		for (size_t i = range.first; i != range.second; ++i) {
			RealNumber v = halfstep * rmasses[i];
			speeds[i].x -= v * grediants[i].x;
			speeds[i].y -= v * grediants[i].y;
			speeds[i].z -= v * grediants[i].z;
		}
		env.wait();
	}

	
	void VelocityVerletControl::act(SimulationEnvironment& env, EnergySet& eset, size_t istep) {
		RealNumber halfstep = timestep / 2;
		RealNumber halfstep2 = halfstep * timestep;
		std::pair<size_t, size_t> range = env.range(eset.nAtoms);
		for (size_t i = range.first; i != range.second; ++i) {
			RealNumber v = halfstep2 * rmasses[i];
			coords[i].x += timestep * speeds[i].x - v * grediants[i].x;
			coords[i].y += timestep * speeds[i].y - v * grediants[i].y;
			coords[i].z += timestep * speeds[i].z - v * grediants[i].z;

			v = halfstep * rmasses[i];
			speeds[i].x -= v * grediants[i].x;
			speeds[i].y -= v * grediants[i].y;
			speeds[i].z -= v * grediants[i].z;
		}
		env.wait();
		eset.getEG(env, energies, grediants);
		for (size_t i = range.first; i != range.second; ++i) {
			RealNumber v = halfstep * rmasses[i];
			speeds[i].x -= v * grediants[i].x;
			speeds[i].y -= v * grediants[i].y;
			speeds[i].z -= v * grediants[i].z;
		}
		env.wait();
	}

	MolecularDynamicJob::MolecularDynamicJob() 
		: md_control(new VelocityVerletControl)
	{ }


	namespace {

		struct SimulationJob_MD
			: public SimulationEnvironment
		{
			void operator()() {
				MolecularDynamicJob* mdjob = dynamic_cast<MolecularDynamicJob*>(job);
				assert (dynamic_cast<MolecularDynamicJob*>(job) != NULL);

				
				efunc->start(*this);
				mdjob->md_control->start(*this, *efunc);
				controls->start(*this, *efunc);

				size_t i = 1;
				for (; !stop() && i <= nstep; ++i) {
					mdjob->md_control->act(*this, *efunc, i);
					controls->act(*this, *efunc, i);
				}
				
				controls->finish(*this, *efunc, i);
				mdjob->md_control->finish(*this, *efunc, i);
				efunc->finish(*this);
			}

			void finalize() {
				controls->finalize(*this, *efunc);
				efunc->finalize(*this);
			}

			EnergySet* efunc;
			ControlSet* controls;
			size_t nstep;
		};
	
	}

	void MolecularDynamicJob::run_md(EnergySet& efunc, ControlSet& controls, size_t nstep, RealNumber timestep) {
		md_control->timestep = timestep;
		stop(false);
		std::vector<SimulationJob_MD> env(getCPU());
		for (size_t i = 0; i < getCPU(); ++i) {
			env[i].efunc = &efunc;
			env[i].controls = &controls;
			env[i].nstep = nstep;
		} 
		runjob(env);
	}

}
